import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import hilbert, butter, filtfilt

# 生成实信号
fs = 1000
t = np.linspace(0, 1, fs, endpoint=False)
FREQ1=50
FREQ2=200
FREQ3=203
signal = np.sin(2 * np.pi *FREQ1 * t) + 0.2 * np.sin(2 * np.pi * FREQ2 * t) + 0.2 *np.sin(2 * np.pi * FREQ3 * t)

# 设计带通滤波器
def butter_bandpass(lowcut, highcut, fs, order=5):
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype='band')
    return b, a

def butter_bandpass_filter(data, lowcut, highcut, fs, order=5):
    b, a = butter_bandpass(lowcut, highcut, fs, order=order)
    y = filtfilt(b, a, data)
    return y

# 实信号带通滤波
lowcut = 180
highcut = 220
filtered_real = butter_bandpass_filter(signal, lowcut, highcut, fs)

# 希尔伯特变换得到解析信号
analytic_signal = hilbert(signal)

# 解析信号带通滤波
filtered_analytic = butter_bandpass_filter(analytic_signal, lowcut, highcut, fs)

# 绘制结果
plt.figure(figsize=(12, 8))

plt.subplot(3, 1, 1)
plt.plot(t, signal, label=f'Original Signal')
plt.title(f'Original Signal(1*{FREQ1}Hz+ 0.2*{FREQ2}Hz + 0.2*{FREQ3}Hz)')
plt.xlabel('Time [s]')
plt.ylabel('Amplitude')
plt.legend()

plt.subplot(3, 1, 2)
plt.plot(t, np.abs(filtered_real), label='Filtered Real Signal')
plt.title(f'BandPassed({lowcut}-{highcut}Hz) - Filtered Real Signal')
plt.xlabel('Time [s]')
plt.ylabel('Amplitude')
plt.legend()

plt.subplot(3, 1, 3)
plt.plot(t, np.abs(filtered_analytic), label='Filtered Analytic Signal (abs)')
plt.title(f'BandPassed({lowcut}-{highcut}Hz) - Filtered Analytic Signal ')
plt.xlabel('Time [s]')
plt.ylabel('Amplitude')
plt.legend()

plt.tight_layout()
plt.show()